/**
 * @file hs_cuda_inputwise.cuh
 * @author iam002
 * @brief 并行应用场景: 输入逐元素运算
 * @version 1.0
 * @date 2024-10-15
 * 
 * @copyright Copyright (c) 2024
 * 
 */
#ifndef __HS_CUDA_INPUTWISE_CUH__
#define __HS_CUDA_INPUTWISE_CUH__

#include "hs_cuda_utils.cuh"

namespace hs
{

namespace cuda
{

namespace inputwise
{

/**************************************************************

inputwise 功能算子与 elemwise 功能算子的区别:

1. 重载()函数的形参不再包含输出数据信息, 重载()函数无返回值;
2. 输出数据信息以指针的形式作为功能算子的结构成员，每次运算都会修改
   这个输出数据指针，因此不能使用 const 修饰重载()函数;
3. 注意输出数据指针必须执行GPU的地址, 不能是主机端!

功能算子演示:

struct HistStatisFunctor
{

    HistStatisFunctor(int bin, float* p_hist)
        : m_bin(bin), m_p_hist(p_hist)
    { }

    __device__ void operator()(float val)
    {
        int rval = (int)rintf(val); // 四舍五入得到索引
        if (rval < m_bin) {
            // 原子加
            atomicAdd(&m_p_hist[rval], 1);
        }
    }

private:
    int m_bin;
    float* m_p_hist; // 输出直返图, 长度为m_bin
};

**************************************************************/


/* 解包, 执行功能单元, 展开可变长输入参数.
 * 功能单元调用: functor(a, b, ...)
 */
template<int pack_size, typename FunctorType, typename... InputType>
__device__ void applyPack(FunctorType functor
    , const Packed<pack_size, InputType>&... src_pack)
{
    #pragma unroll
    for (int i = 0; i < pack_size; ++i) {
        functor((src_pack.elem[i])...);
    }
}

/* 核函数, 请注意, 核函数形参涉及到数组和指针的传递时, 应该保证它们是在device端, 这真的很重要;
 * 另外, 核函数也不要使用引用.
 */
template<int pack_size, typename FunctorType, typename... InputType>
__global__ void __launch_bounds__(KERNEL_BLOCK_SIZE)
applyGeneric(FunctorType functor
        , int num_pack
        , const Packed<pack_size, InputType>*... src_pack
        , int num_tail
        , const InputType*... p_in_tail)
{
    const int global_tid = threadIdx.x + blockIdx.x * blockDim.x;
    for (int i = global_tid; i < num_pack; i += blockDim.x * gridDim.x) {
        /* 处理包 */
        applyPack<pack_size, FunctorType, InputType...>(
            functor, (src_pack[i])...
        );
    }
    if (global_tid < num_tail) {
        /* 处理尾部数据 */
        functor((p_in_tail[global_tid])...);
    }
}


/* 准备工作, 用于启动核函数 */
template<size_t pack_size, typename FunctorType, typename... InputType>
cudaError_t launchKernel(FunctorType functor, int N, const InputType*... p_in)
{
    /* 计算打包数和尾部剩余数据个数 */
    const int num_pack = N / pack_size;
    const int tail_offset = num_pack * pack_size;
    const int num_tail = N - tail_offset;

    /* 计算线程块数 */
    int num_blocks;
    {
        // 不调用 hs_cuda_utils.cu
        cudaError_t err = GetNumBlocks(num_pack, &num_blocks);
        if (err != cudaSuccess) {
            return err;
        }
    }

    /* 调用核函数 */
    applyGeneric<pack_size, FunctorType, InputType...><<<num_blocks, KERNEL_BLOCK_SIZE, 0>>>(
        functor
        , num_pack
        , (reinterpret_cast<const Packed<pack_size, InputType>*>(p_in))...
        , num_tail
        , (p_in + tail_offset)... 
    );

    return cudaPeekAtLastError();
}


/* 再次封装, 隐藏 pack_size */
template<typename FunctorType, typename... InputType>
struct GenericLauncher
{
    static cudaError_t launch(FunctorType functor, int N, const InputType*... p_in)
    {
        constexpr int pack_size = PackSize<InputType...>();
        
        cudaError_t err;
        if (isAlignedForPack<pack_size, InputType...>(p_in...)) {
            err = launchKernel<pack_size, FunctorType, InputType...>(functor, N, p_in...);
        }
        else {
            err = launchKernel<1, FunctorType, InputType...>(functor, N, p_in...);
        }
        return err;
    }
};


/**
 * @brief 一元运算模板
 * 
 * @tparam FunctorType 重载()运算的结构体
 * @tparam InputType 输入数据类型
 * @param N 数组长度
 * @param p_in 输入数据指针
 * @param functor 函数器
 * @return cudaError_t 
 */
template<typename FunctorType, typename InputType>
inline cudaError_t Unary(int N, const InputType* p_in, FunctorType functor)
{  
    return GenericLauncher<FunctorType, InputType>::launch(functor, N, p_in);
}


/**
 * @brief 二元运算模板
 * 
 * @tparam FunctorType 重载()运算的结构体
 * @tparam InputTypeA  输入数据类型A
 * @tparam InputTypeB  输入数据类型B
 * @param N 数组长度
 * @param p_in_a 输入数据指针A
 * @param p_in_b 输入数据指针B
 * @param functor 函数器
 * @return cudaError_t 
 */
template<typename FunctorType, typename InputTypeA, typename InputTypeB>
inline cudaError_t Binary(int N, const InputTypeA* p_in_a, const InputTypeB* p_in_b, const FunctorType & functor)
{
    return GenericLauncher<FunctorType, InputTypeA, InputTypeB>::launch(functor, N, p_in_a, p_in_b);
}


/**
 * @brief 三元运算模板
 * 
 * @tparam FunctorType  重载()运算的结构体
 * @tparam InputTypeA   输入数据类型A
 * @tparam InputTypeB   输入数据类型B
 * @tparam InputTypeC   输入数据类型C
 * @param N 数组长度
 * @param p_in_a 输入数据指针A
 * @param p_in_b 输入数据指针B
 * @param p_in_c 输入数据指针C
 * @param functor 函数器
 * @return cudaError_t 
 */
template<typename FunctorType, typename InputTypeA, typename InputTypeB, typename InputTypeC>
inline cudaError_t Ternary(int N, const InputTypeA* p_in_a, const InputTypeB* p_in_b, const InputTypeC* p_in_c, const FunctorType & functor)
{
    return GenericLauncher<FunctorType, InputTypeA, InputTypeB, InputTypeC>::launch(functor, N, p_in_a, p_in_b, p_in_c);
}

} // end namespace inputwise

namespace inputwiseIndex
{

/**************************************************************

功能算子演示:

struct PaddingZeroFunctor
{

    PaddingZeroFunctor(int dst_size_x, int dst_size_y, float* p_dst_data, int src_size_x, int src_size_y)
        : m_dst_size_x(dst_size_x), m_dst_size_y(dst_size_y)
        , m_p_dst_data(p_dst_data)
        , m_src_size_x(src_size_x), m_src_size_y(src_size_y)
    { }

    __device__ void operator()(int id, float val)
    {
        if (id < m_src_size_x * m_src_size_y) {
            // id 是输入图像的线性索引, 转换为平面索引
            int yid = id / m_src_size_x;
            int xid = id - yid * m_src_size_x;

            // 加上偏移量
            yid = yid + 1;
            xid = xid + 1;

            // 转换为输出图线性索引
            int dst_id = xid + yid * m_dst_size_x;
            m_p_dst_data[dst_id] = val;
        }
    }

private:
    int m_src_size_y;
    int m_src_size_x;
    int m_dst_size_x;
    int m_dst_size_y;
    float* m_p_dst_data;
};

**************************************************************/

/* 解包, 执行功能单元, 展开可变长输入参数.
 * 功能单元调用: functor(ID, a, b, ...), ID为输入索引
 */
template<int pack_size, typename FunctorType, typename... InputType>
__device__ void applyPack(FunctorType functor
    , int pack_id
    , const Packed<pack_size, InputType>&... src_pack)
{
    #pragma unroll
    for (int i = 0; i < pack_size; ++i) {
        functor(pack_id * pack_size + i, (src_pack.elem[i])...);
    }
}

/* 核函数, 请注意, 核函数形参涉及到数组和指针的传递时, 应该保证它们是在device端, 这真的很重要;
 * 另外, 核函数也不要使用引用.
 */
template<int pack_size, typename FunctorType, typename... InputType>
__global__ void __launch_bounds__(KERNEL_BLOCK_SIZE)
applyGeneric(FunctorType functor
        , int num_pack
        , const Packed<pack_size, InputType>*... src_pack
        , int num_tail
        , const InputType*... p_in_tail)
{
    const int global_tid = threadIdx.x + blockIdx.x * blockDim.x;
    for (int i = global_tid; i < num_pack; i += blockDim.x * gridDim.x) {
        /* 处理包 */
        applyPack<pack_size, FunctorType, InputType...>(
            functor, i, (src_pack[i])...
        );
    }
    if (global_tid < num_tail) {
        /* 处理尾部数据 */
        const int tail_offset = num_pack * pack_size;
        functor(global_tid + tail_offset, (p_in_tail[global_tid])...);
    }
}


/* 准备工作, 用于启动核函数 */
template<size_t pack_size, typename FunctorType, typename... InputType>
cudaError_t launchKernel(FunctorType functor, int N, const InputType*... p_in)
{
    /* 计算打包数和尾部剩余数据个数 */
    const int num_pack = N / pack_size;
    const int tail_offset = num_pack * pack_size;
    const int num_tail = N - tail_offset;

    /* 计算线程块数 */
    int num_blocks;
    {
        // 不调用 hs_cuda_utils.cu
        cudaError_t err = GetNumBlocks(num_pack, &num_blocks);
        if (err != cudaSuccess) {
            return err;
        }
    }

    /* 调用核函数 */
    applyGeneric<pack_size, FunctorType, InputType...><<<num_blocks, KERNEL_BLOCK_SIZE, 0>>>(
        functor
        , num_pack
        , (reinterpret_cast<const Packed<pack_size, InputType>*>(p_in))...
        , num_tail
        , (p_in + tail_offset)... 
    );

    return cudaPeekAtLastError();
}


/* 再次封装, 隐藏 pack_size */
template<typename FunctorType, typename... InputType>
struct GenericLauncher
{
    static cudaError_t launch(FunctorType functor, int N, const InputType*... p_in)
    {
        constexpr int pack_size = PackSize<InputType...>();
        
        cudaError_t err;
        if (isAlignedForPack<pack_size, InputType...>(p_in...)) {
            err = launchKernel<pack_size, FunctorType, InputType...>(functor, N, p_in...);
        }
        else {
            err = launchKernel<1, FunctorType, InputType...>(functor, N, p_in...);
        }
        return err;
    }
};


/**
 * @brief 一元运算模板
 * 
 * @tparam FunctorType 重载()运算的结构体
 * @tparam InputType 输入数据类型
 * @param N 数组长度
 * @param p_in 输入数据指针
 * @param functor 函数器
 * @return cudaError_t 
 */
template<typename FunctorType, typename InputType>
inline cudaError_t Unary(int N, const InputType* p_in, FunctorType functor)
{  
    return GenericLauncher<FunctorType, InputType>::launch(functor, N, p_in);
}


/**
 * @brief 二元运算模板
 * 
 * @tparam FunctorType 重载()运算的结构体
 * @tparam InputTypeA  输入数据类型A
 * @tparam InputTypeB  输入数据类型B
 * @param N 数组长度
 * @param p_in_a 输入数据指针A
 * @param p_in_b 输入数据指针B
 * @param functor 函数器
 * @return cudaError_t 
 */
template<typename FunctorType, typename InputTypeA, typename InputTypeB>
inline cudaError_t Binary(int N, const InputTypeA* p_in_a, const InputTypeB* p_in_b, const FunctorType & functor)
{
    return GenericLauncher<FunctorType, InputTypeA, InputTypeB>::launch(functor, N, p_in_a, p_in_b);
}


/**
 * @brief 三元运算模板
 * 
 * @tparam FunctorType  重载()运算的结构体
 * @tparam InputTypeA   输入数据类型A
 * @tparam InputTypeB   输入数据类型B
 * @tparam InputTypeC   输入数据类型C
 * @param N 数组长度
 * @param p_in_a 输入数据指针A
 * @param p_in_b 输入数据指针B
 * @param p_in_c 输入数据指针C
 * @param functor 函数器
 * @return cudaError_t 
 */
template<typename FunctorType, typename InputTypeA, typename InputTypeB, typename InputTypeC>
inline cudaError_t Ternary(int N, const InputTypeA* p_in_a, const InputTypeB* p_in_b, const InputTypeC* p_in_c, const FunctorType & functor)
{
    return GenericLauncher<FunctorType, InputTypeA, InputTypeB, InputTypeC>::launch(functor, N, p_in_a, p_in_b, p_in_c);
}


} // end namespace inputwiseIndex

} // end namespace cuda

} // end namespace hs

#endif
